from collections import deque

from nanovllm.config import Config
from nanovllm.engine.sequence import Sequence, SequenceStatus
from nanovllm.engine.block_manager import BlockManager


class Scheduler:
    """
    这段代码定义了一个 Scheduler（调度器）类，主要用于在大语言模型（LLM）推理过程中管理序列（Sequence）的生命周期和资源分配，
    核心功能是协调 “等待中” 和 “运行中” 的序列，合理分配 KV 缓存资源，并决定下一批需要处理的序列。

    在 LLM 推理中，“序列（Sequence）” 可以理解为一个完整的生成任务（例如一次对话、一段文本续写）。
    每个序列在生成过程中需要消耗计算资源（如 GPU 算力）和存储资源（如 KV 缓存，用于存储中间注意力计算结果）。
    Scheduler 的作用就是高效调度这些序列，确保资源利用率最大化，同时满足并发和性能需求。

    Sequence：代表一个推理任务（如一次文本生成），包含生成的 token、状态、缓存需求等信息。
    SequenceStatus：序列的状态枚举（如 WAITING 等待中、RUNNING 运行中、FINISHED 已完成）。
    BlockManager：管理 KV 缓存块的分配与释放（KV 缓存是 LLM 推理中存储中间结果的关键资源，直接影响性能）。
    双端队列：waiting（等待处理的序列）和 running（正在处理的序列），用于高效管理序列的进出。
    """

    def __init__(self, config: Config):
        self.max_num_seqs = config.max_num_seqs # 最大并发序列数（系统支持的同时处理的任务数）
        self.max_num_batched_tokens = config.max_num_batched_tokens  # 每批最大处理的token数
        self.eos = config.eos  # EOS标记（生成结束的标志，如<|endoftext|>）
        self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size) # KV缓存管理器
        self.waiting: deque[Sequence] = deque()  # 等待队列（新任务或被抢占的任务）
        self.running: deque[Sequence] = deque() # 运行队列（正在处理的任务）

    def is_finished(self):
        return not self.waiting and not self.running # 等待队列和运行队列都为空时，返回True

    def add(self, seq: Sequence):
        self.waiting.append(seq) # 新任务加入等待队列，等待调度

    def schedule(self) -> tuple[list[Sequence], bool]:
        """
        核心调度逻辑（决定下一批处理的序列）
        Prefill 阶段：处理序列的初始输入（如用户的提问文本），需要一次性计算所有输入 token 的 KV 缓存，消耗资源较多。
        Decode 阶段：在已有 KV 缓存的基础上，逐个生成后续 token（如模型的回答），每次只处理 1 个新 token，资源消耗较少。
        :return: tuple
        """
        # prefill
        scheduled_seqs = [] # 本次要处理的序列
        num_seqs = 0 # 已调度的序列数
        num_batched_tokens = 0 # 已调度的token总数
        while self.waiting and num_seqs < self.max_num_seqs:
            seq = self.waiting[0]  # 取等待队列的第一个序列
            # 检查是否超过资源限制：token数超限 或 KV缓存块不足
            if num_batched_tokens + len(seq) > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq):
                break
            # 资源足够，调度该序列
            num_seqs += 1
            self.block_manager.allocate(seq) # 分配KV缓存块
            # 累加新增的token数（总长度 - 已缓存的token数，避免重复计算）
            num_batched_tokens += len(seq) - seq.num_cached_tokens
            seq.status = SequenceStatus.RUNNING  # 更新状态为运行中
            self.waiting.popleft() # 从等待队列移除
            self.running.append(seq) # 加入运行队列
            scheduled_seqs.append(seq)  # 加入本次调度列表
        if scheduled_seqs:
            return scheduled_seqs, True # 返回调度的序列，True表示是prefill阶段
        # 优先处理等待队列中的新序列，尽可能填满一批（不超过最大序列数和最大 token 数）。
        # 必须通过 BlockManager 检查并分配 KV 缓存块（资源不足则停止调度）。

        # decode
        # 如果 Prefill 阶段没有调度任何序列（等待队列空或资源不足），则进入 Decode 阶段，处理运行中的序列：
        while self.running and num_seqs < self.max_num_seqs:
            seq = self.running.popleft()  # 取运行队列的第一个序列
            while not self.block_manager.can_append(seq): # 检查是否可以为该序列追加新token（需要额外的KV缓存块）
                if self.running: # 缓存不足，需要抢占其他序列的资源
                    self.preempt(self.running.pop())   # 抢占运行队列最后一个序列
                else:
                    self.preempt(seq) # 没有其他序列可抢占，只能抢占当前序列
                    break
            else:
                num_seqs += 1 # 可以追加token，调度该序列
                self.block_manager.may_append(seq) # 预留缓存块（准备生成新token）
                scheduled_seqs.append(seq)
        assert scheduled_seqs # 确保至少调度了一个序列
        self.running.extendleft(reversed(scheduled_seqs)) # 将调度的序列放回运行队列（保持顺序）
        # 处理运行中的序列，每次生成 1 个新 token，需要检查 KV 缓存是否足够追加
        # 若缓存不足，通过 preempt 方法抢占其他序列的资源（将其移回等待队列，释放缓存）。
        # 调度完成后，将序列放回运行队列（extendleft 保证顺序不变）。
        return scheduled_seqs, False # 返回调度的序列，False表示是decode阶段

    def preempt(self, seq: Sequence):
        """ 资源抢占（缓存不足时的处理） """
        seq.status = SequenceStatus.WAITING # 状态改为等待中
        self.block_manager.deallocate(seq) # 释放该序列的KV缓存块
        self.waiting.appendleft(seq) # 加入等待队列的头部（优先再次被调度）

    def postprocess(self, seqs: list[Sequence], token_ids: list[int]) -> list[bool]:
        """ 处理生成结果，更新序列状态 """
        for seq, token_id in zip(seqs, token_ids):
            seq.append_token(token_id) # 为序列追加新生成的token
            # 检查是否结束：生成了EOS token 或 达到最大生成长度
            if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
                seq.status = SequenceStatus.FINISHED # 标记为已完成
                self.block_manager.deallocate(seq)  # 释放KV缓存块
                self.running.remove(seq) # 从运行队列移除

"""
cheduler 的核心目标是高效利用资源（KV 缓存、计算能力），最大化并发处理能力，其工作流程可概括为：

新序列通过 add 方法进入等待队列。
schedule 方法优先处理 Prefill 阶段的序列（初始输入），分配足够的 KV 缓存。
若 Prefill 无序列可调度，则处理 Decode 阶段的序列（生成后续 token），缓存不足时通过抢占机制释放资源。
生成结束后，postprocess 方法更新序列状态，释放资源。

这种设计确保了 LLM 推理过程中资源的高效利用，同时支持高并发的生成任务（如同时处理多个用户的提问）。
"""